package org.apache.flink.runtime.resourceprovider;

import java.io.File;
import java.io.IOException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;

import org.apache.flink.configuration.CoreOptions;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.blob.PermanentBlobService;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.entrypoint.WorkingDirectory;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.memory.SharedResources;
import org.apache.flink.runtime.providermanager.ConsumingItem;
import org.apache.flink.runtime.providermanager.UnresolvedProviderManagerLocation;
import org.apache.flink.runtime.resourceprovider.slot.DefaultTimerService;
import org.apache.flink.runtime.resourceprovider.slot.FileSlotAllocationSnapshotPersistenceService;
import org.apache.flink.runtime.resourceprovider.slot.NoOpSlotAllocationSnapshotPersistenceService;
import org.apache.flink.runtime.resourceprovider.slot.SlotAllocationSnapshotPersistenceService;
import org.apache.flink.runtime.resourceprovider.slot.TaskSlotTable;
import org.apache.flink.runtime.resourceprovider.slot.TaskSlotTableImpl;
import org.apache.flink.runtime.resourceprovider.slot.TimerService;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.ERException;
import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * @类名: ResourceProviderServices
 * @说明: STO，资源提供者服务和组件容器
 *
 * @author   leehom
 * @Date	 2024年1月26日 下午2:38:25
 * 修改记录：
 *
 * @see 	 
 */
/**
 * Container for {@link ResourceProvider} services such as the {@link MemoryManager}, {@link IOManager},
 * {@link ShuffleEnvironment}. All services are exclusive to a single {@link ResourceProvider}.
 * Consequently, the respective {@link ResourceProvider} is responsible for closing them.
 */
public class ResourceProviderServices {
    private static final Logger LOG = LoggerFactory.getLogger(ResourceProviderServices.class);

    /** ProviderManager services. */
    private final UnresolvedProviderManagerLocation unresolvedTaskManagerLocation;

    private final long managedMemorySize;
    private final TaskSlotTable<ConsumingItem> taskSlotTable;
    private final ConsumerTable consumerTable;
    private final ConsumerLeaderService jobLeaderService;
    private final ExecutorService ioExecutor;
    private final SlotAllocationSnapshotPersistenceService slotAllocationSnapshotPersistenceService;
    private final SharedResources sharedResources;

    ResourceProviderServices(
            UnresolvedProviderManagerLocation unresolvedTaskManagerLocation,
            long managedMemorySize,
            TaskSlotTable<ConsumingItem> taskSlotTable,
            ConsumerTable consumerTable,
            ConsumerLeaderService jobLeaderService,
            ExecutorService ioExecutor,
            SlotAllocationSnapshotPersistenceService slotAllocationSnapshotPersistenceService,
            SharedResources sharedResources) {

        this.unresolvedTaskManagerLocation =
                Preconditions.checkNotNull(unresolvedTaskManagerLocation);
        this.managedMemorySize = managedMemorySize;
        this.taskSlotTable = Preconditions.checkNotNull(taskSlotTable);
        this.consumerTable = Preconditions.checkNotNull(consumerTable);
        this.jobLeaderService = Preconditions.checkNotNull(jobLeaderService);
        this.ioExecutor = Preconditions.checkNotNull(ioExecutor);
        this.slotAllocationSnapshotPersistenceService = slotAllocationSnapshotPersistenceService;
        this.sharedResources = Preconditions.checkNotNull(sharedResources);
    }

    // --------------------------------------------------------------------------------------------
    //  Getter/Setter
    // --------------------------------------------------------------------------------------------

    public long getManagedMemorySize() {
        return managedMemorySize;
    }

    public UnresolvedProviderManagerLocation getUnresolvedTaskManagerLocation() {
        return unresolvedTaskManagerLocation;
    }

    public TaskSlotTable<ConsumingItem> getTaskSlotTable() {
        return taskSlotTable;
    }

    public ConsumerTable getConsumerTable() {
        return consumerTable;
    }

    public ConsumerLeaderService getJobLeaderService() {
        return jobLeaderService;
    }

    public Executor getIOExecutor() {
        return ioExecutor;
    }

    public SharedResources getSharedResources() {
        return sharedResources;
    }

    // --------------------------------------------------------------------------------------------
    //  Shut down method
    // --------------------------------------------------------------------------------------------

    /** Shuts the {@link ResourceProvider} services down. */
    public void shutDown() throws ERException {

        Exception exception = null;

        try {
            taskSlotTable.close();
        } catch (Exception e) {
            exception = e;
        }

        try {
            jobLeaderService.stop();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }

        try {
            ioExecutor.shutdown();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }

        try {
            consumerTable.close();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }

        if (exception != null) {
            throw new ERException(
                    "Could not properly shut down the ProviderManager services.", exception);
        }
    }

    // --------------------------------------------------------------------------------------------
    //  Static factory methods for task manager services
    // --------------------------------------------------------------------------------------------

    /**
     * Creates and returns the task manager services.
     *
     * @param resourceProviderServicesConfiguration task manager configuration
     * @param permanentBlobService permanentBlobService used by the services
     * @param taskManagerMetricGroup metric group of the task manager
     * @param ioExecutor executor for async IO operations
     * @param fatalErrorHandler to handle class loading OOMs
     * @param workingDirectory the working directory of the process
     * @return task manager components
     * @throws Exception
     */
    public static ResourceProviderServices fromConfiguration(
            ResourceProviderServicesConfiguration resourceProviderServicesConfiguration,
            PermanentBlobService permanentBlobService,
            MetricGroup taskManagerMetricGroup,
            ExecutorService ioExecutor,
            FatalErrorHandler fatalErrorHandler,
            WorkingDirectory workingDirectory)
            throws Exception {

        // pre-start checks
        checkTempDirs(resourceProviderServicesConfiguration.getTmpDirPaths());

        final UnresolvedProviderManagerLocation unresolvedTaskManagerLocation =
                new UnresolvedProviderManagerLocation(
                        resourceProviderServicesConfiguration.getResourceID(),
                        resourceProviderServicesConfiguration.getExternalAddress(),
                        resourceProviderServicesConfiguration.getNodeId());

        final TaskSlotTable<ConsumingItem> taskSlotTable =
                createTaskSlotTable(
                        resourceProviderServicesConfiguration.getNumberOfSlots(),
                        resourceProviderServicesConfiguration.getTaskExecutorResourceSpec(),
                        resourceProviderServicesConfiguration.getTimerServiceShutdownTimeout(),
                        resourceProviderServicesConfiguration.getPageSize(),
                        ioExecutor);

        final ConsumerTable jobTable = DefaultConsumerTable.create();

        final ConsumerLeaderService jobLeaderService =
                new DefaultConsumerLeaderService(
                        unresolvedTaskManagerLocation,
                        resourceProviderServicesConfiguration.getRetryingRegistrationConfiguration());

        final boolean failOnJvmMetaspaceOomError =
                resourceProviderServicesConfiguration
                        .getConfiguration()
                        .getBoolean(CoreOptions.FAIL_ON_USER_CLASS_LOADING_METASPACE_OOM);
        final boolean checkClassLoaderLeak =
                resourceProviderServicesConfiguration
                        .getConfiguration()
                        .getBoolean(CoreOptions.CHECK_LEAKED_CLASSLOADER);

        final SlotAllocationSnapshotPersistenceService slotAllocationSnapshotPersistenceService;

        if (resourceProviderServicesConfiguration.isLocalRecoveryEnabled()) {
            slotAllocationSnapshotPersistenceService =
                    new FileSlotAllocationSnapshotPersistenceService(
                            workingDirectory.getSlotAllocationSnapshotDirectory());
        } else {
            slotAllocationSnapshotPersistenceService =
                    NoOpSlotAllocationSnapshotPersistenceService.INSTANCE;
        }

        return new ResourceProviderServices(
                unresolvedTaskManagerLocation,
                resourceProviderServicesConfiguration.getManagedMemorySize().getBytes(),
                taskSlotTable,
                jobTable,
                jobLeaderService,
                ioExecutor,
                slotAllocationSnapshotPersistenceService,
                new SharedResources());
    }

    private static TaskSlotTable<ConsumingItem> createTaskSlotTable(
            final int numberOfSlots,
            final ResourceProviderResourceSpec taskExecutorResourceSpec,
            final long timerServiceShutdownTimeout,
            final int pageSize,
            final Executor memoryVerificationExecutor) {
        final TimerService<AllocationID> timerService =
                new DefaultTimerService<>(
                        new ScheduledThreadPoolExecutor(1), timerServiceShutdownTimeout);
        return new TaskSlotTableImpl<>(
                numberOfSlots,
                ResourceProviderResourceUtils.generateTotalAvailableResourceProfile(
                        taskExecutorResourceSpec),
                ResourceProviderResourceUtils.generateDefaultSlotResourceProfile(
                        taskExecutorResourceSpec, numberOfSlots),
                pageSize,
                timerService,
                memoryVerificationExecutor);
    }

    /**
     * Validates that all the directories denoted by the strings do actually exist or can be
     * created, are proper directories (not files), and are writable.
     *
     * @param tmpDirs The array of directory paths to check.
     * @throws IOException Thrown if any of the directories does not exist and cannot be created or
     *     is not writable or is a file, rather than a directory.
     */
    private static void checkTempDirs(String[] tmpDirs) throws IOException {
        for (String dir : tmpDirs) {
            if (dir != null && !dir.equals("")) {
                File file = new File(dir);
                if (!file.exists()) {
                    if (!file.mkdirs()) {
                        throw new IOException(
                                "Temporary file directory "
                                        + file.getAbsolutePath()
                                        + " does not exist and could not be created.");
                    }
                }
                if (!file.isDirectory()) {
                    throw new IOException(
                            "Temporary file directory "
                                    + file.getAbsolutePath()
                                    + " is not a directory.");
                }
                if (!file.canWrite()) {
                    throw new IOException(
                            "Temporary file directory "
                                    + file.getAbsolutePath()
                                    + " is not writable.");
                }

                if (LOG.isInfoEnabled()) {
                    long totalSpaceGb = file.getTotalSpace() >> 30;
                    long usableSpaceGb = file.getUsableSpace() >> 30;
                    double usablePercentage = (double) usableSpaceGb / totalSpaceGb * 100;
                    String path = file.getAbsolutePath();
                    LOG.info(
                            String.format(
                                    "Temporary file directory '%s': total %d GB, "
                                            + "usable %d GB (%.2f%% usable)",
                                    path, totalSpaceGb, usableSpaceGb, usablePercentage));
                }
            } else {
                throw new IllegalArgumentException("Temporary file directory #$id is null.");
            }
        }
    }

    public SlotAllocationSnapshotPersistenceService getSlotAllocationSnapshotPersistenceService() {
        return slotAllocationSnapshotPersistenceService;
    }
}
